#include <hvrl/opengl/GLSLShader.hpp>
#include <hvrl/common/Common.hpp>

#include <opencv2/opencv.hpp>
#include <GL/freeglut.h>

// Include for GLM
#include "glm/gtc/matrix_transform.hpp"
#include "glm/gtx/transform.hpp"
#include "glm/gtx/transform2.hpp"
#include "glm/gtc/type_ptr.hpp"

// Include for CUDA
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <cuda_gl_interop.h>
#include <hvrl/gpgpu/cuda/Prototypes.hpp>


using namespace hvrl;

/// Video camera
cv::Size windowsize(800,600);
cv::VideoCapture cap;
cv::Mat videoframe;

/// Size of the chessboard pattern
static const cv::Size pattern = cv::Size(7,4);
static float squareSize = 0.03f;
static const cv::TermCriteria criteria = cv::TermCriteria(CV_TERMCRIT_ITER|CV_TERMCRIT_EPS, 10, 0.01f);
static const unsigned int nbmaxcaptures = 15;

/// program states
bool undistortcomputed = false;
bool isrunning = true;
bool useundistort = false;
static bool startCalibration = false;

/// timers
static int starttimedisplay = 0;
static int starttimecapture = 0;
static int starttimecalibration = 0;

/// Pixel Buffer Object
GLuint pbo = 0;

/// Texture
GLuint cameratextureid = 0;

/// cuda warpper
struct cudaGraphicsResource *cuda_pbo_resource=0;

/// GLSL program object
GLuint cameraso = 0;

/// Vertex Array Object
struct VertexArrayObject
{
    GLuint vertexarrayid;
    GLuint verticesid;
    GLuint normalsid;
    GLuint texturesid;
    GLuint indicesid;

    GLsizei sizeindices;

    GLint verticesattribid;
    GLint normalsattribid;
    GLint texturesattribid;

};
VertexArrayObject cameraQuad;

/// Transformation matrices
glm::mat4 projection = glm::mat4(1.0);
glm::mat4 view = glm::mat4(1.0);
glm::mat4 mvp = glm::mat4(1.0);



namespace hvrl
{
namespace cuda
{
void runRadialDistortion(unsigned int*);
void updateRadialDistortion(unsigned char*);
void destroyRadialDistortion(void);
bool initRadialDistortion(const float*,const unsigned int&,const unsigned int&);
}
}
bool createFullScreenQuad(const unsigned int& width, const unsigned int& height, VertexArrayObject& vao);
void drawVAO(const VertexArrayObject& vao);
void deleteVAO(VertexArrayObject& vao);


void display(void)
{
    if(undistortcomputed && useundistort)
    {
        unsigned int *d_result;
        size_t num_bytes;
        /// Map the memory with cuda
        cuda::checkStatus(cudaGraphicsMapResources(1, &cuda_pbo_resource, 0));
        cuda::checkStatus(cudaGraphicsResourceGetMappedPointer((void **)&d_result, &num_bytes,cuda_pbo_resource));
        cuda::runRadialDistortion(d_result);
        cuda::checkStatus(cudaGraphicsUnmapResources(1, &cuda_pbo_resource, 0));

        /// Update the texture
        glBindBufferARB(GL_PIXEL_UNPACK_BUFFER_ARB, pbo);
        glBindTexture(GL_TEXTURE_2D, cameratextureid);
        glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, windowsize.width, windowsize.height, GL_BGRA, GL_UNSIGNED_BYTE, 0);
        glBindBufferARB(GL_PIXEL_UNPACK_BUFFER_ARB, 0);
    }

    glClear(GL_COLOR_BUFFER_BIT);
    glUseProgram(cameraso);
    glBindTexture(GL_TEXTURE_2D, cameratextureid);

    GLint loc = glGetUniformLocation(cameraso,"texture");
    if(loc < 0)
        hvrl::Log::add().error("display","glsl location error texture");
    glUniform1i(loc,0);

    loc = glGetUniformLocation(cameraso,"mvpMatrixIn");
    if(loc < 0)
        hvrl::Log::add().error("display","glsl location error mvpMatrixIn");
    glUniformMatrix4fv(loc, 1, GL_FALSE, glm::value_ptr(mvp));
    drawVAO(cameraQuad);
    glBindTexture(GL_TEXTURE_2D,0);
    glUseProgram(0);

    glutSwapBuffers();
#ifdef DEBUG
    GLenum err = glGetError();
    if (err!=GL_NO_ERROR)
    {
        hvrl::Log::add().error("display","GL error: " + std::string((char*)gluErrorString(err)));
                               exit(0);
                           }
#endif
                       }



                       void keyboard(unsigned char key, int x, int y)
{

    switch(key)
    {
    case 27:
    {
        isrunning = false;
        break;
    }
    case 'u':
    {
        startCalibration = true;
        starttimecalibration = glutGet(GLUT_ELAPSED_TIME);
        break;
    }
    case 'r':
    {
        useundistort = !useundistort;
        break;
    }
    case 'h':
    {
        hvrl::Log::add().info("Program usage:");
        hvrl::Log::add().info("\'r\' : activate/deactivate undistortion");
        hvrl::Log::add().info("\'u\' : calibrate the undistortion (require chessboard pattern)");
        hvrl::Log::add().info("\'Esc\' : quit");
    }


    }

}

void idle(void)
{

    /*
        Update the camera
    */
    int diff = glutGet(GLUT_ELAPSED_TIME) - starttimecapture;
    if(diff > 33)
    {

        cap >> videoframe;
        starttimecapture = glutGet(GLUT_ELAPSED_TIME);
    }

    if(startCalibration)
    {
        static std::vector<std::vector<cv::Point2f> > points;
        static std::vector<std::vector<cv::Point3f> > corners(1);
        static unsigned int currentcaptures = 0;
        if(currentcaptures<nbmaxcaptures)
        {
            diff = glutGet(GLUT_ELAPSED_TIME) - starttimecalibration;
            std::vector<cv::Point2f> cornerspoints;
            bool found = cv::findChessboardCorners(videoframe, pattern, cornerspoints, CV_CALIB_CB_FAST_CHECK|CV_CALIB_CB_ADAPTIVE_THRESH|CV_CALIB_CB_NORMALIZE_IMAGE);
            cv::drawChessboardCorners(videoframe,pattern,cornerspoints,found);
            if(found && diff > 300)
            {
                cv::Mat gray;
                cvtColor(videoframe, gray, CV_RGB2GRAY);
                cv::cornerSubPix(gray, cornerspoints, cv::Size(11,11), cv::Size(-1,-1), criteria);
                points.push_back(cornerspoints);
                ++currentcaptures;
                hvrl::Log::add().info("Capture "+hvrl::numberToString(currentcaptures)+"/"+hvrl::numberToString(nbmaxcaptures));
                starttimecalibration = glutGet(GLUT_ELAPSED_TIME);
            }
        }
        else
        {
            cv::Mat_<float> K;
            cv::Mat distCoeffs;
            std::vector<cv::Mat> rvecs;
            std::vector<cv::Mat> tvecs;
            K = cv::Mat::eye(3, 3, CV_64F);
            distCoeffs = cv::Mat::zeros(8, 1, CV_32F);
            corners[0].resize(0);
            for( int i = 0; i < pattern.height; i++ )
                for( int j = 0; j < pattern.width; j++ )
                    corners[0].push_back(cv::Point3f(float(j*squareSize), float(i*squareSize), 0));
            corners.resize(points.size(),corners[0]);

            calibrateCamera(corners, points, videoframe.size(), K, distCoeffs, rvecs, tvecs, CV_CALIB_FIX_K4|CV_CALIB_FIX_K5);

            points.clear();
            rvecs.clear();
            tvecs.clear();
            corners.clear();
            cv::Mat map1, map2;
            cv::initUndistortRectifyMap(K, distCoeffs, cv::Mat(), K, videoframe.size(), CV_32FC1, map1, map2);
            float* map = new float[2*videoframe.rows*videoframe.cols];

            for(int i = 0; i < videoframe.rows; ++i)
            {
                for(int j = 0; j < videoframe.cols; ++j)
                {
                    map[(i*videoframe.size().width + j)*2] = map1.at<float>(i,j);
                    map[(i*videoframe.size().width + j)*2+1] = map2.at<float>(i,j);
                }
            }

            cuda::initRadialDistortion(map,videoframe.cols,videoframe.rows);

            cuda::updateRadialDistortion(videoframe.data);

            delete[] map;
            startCalibration = false;
            undistortcomputed = true;
            useundistort = true;
        }


    }

    /*
        Update the display
    */
    diff = glutGet(GLUT_ELAPSED_TIME) - starttimedisplay;
    if(diff > 33)
    {
        if(undistortcomputed && useundistort)
        {
            cuda::updateRadialDistortion(videoframe.data);
        }
        else
        {
            glBindBufferARB(GL_PIXEL_UNPACK_BUFFER_ARB, pbo);
            GLubyte* ptr = (GLubyte*)glMapBufferARB(GL_PIXEL_UNPACK_BUFFER_ARB, GL_WRITE_ONLY_ARB);
            if(ptr)
            {
                // update data directly on the mapped buffer
                memcpy(ptr,videoframe.data,videoframe.rows*videoframe.cols*3*sizeof(unsigned char));
                glUnmapBufferARB(GL_PIXEL_UNPACK_BUFFER_ARB); // release the mapped buffer
            }

            glBindTexture(GL_TEXTURE_2D, cameratextureid);
            glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, windowsize.width, windowsize.height, GL_BGR, GL_UNSIGNED_BYTE, 0);
            glBindBufferARB(GL_PIXEL_UNPACK_BUFFER_ARB, 0);
        }

        glutPostRedisplay();
        starttimedisplay = glutGet(GLUT_ELAPSED_TIME);
    }


}

bool initCuda()
{
    cap >> videoframe;

    glGenBuffers(1, &pbo);
    if(pbo==0)
    {
        std::cerr << "Failed to create a pbo" << std::endl;
        return false;
    }
    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, pbo);
    glBufferData(GL_PIXEL_UNPACK_BUFFER, videoframe.rows*videoframe.cols*sizeof(GLubyte)*4, videoframe.data, GL_STREAM_DRAW_ARB);
    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0);
    cuda::checkStatus(cudaGraphicsGLRegisterBuffer(&cuda_pbo_resource, pbo,cudaGraphicsMapFlagsWriteDiscard));

    glGenTextures(1,&cameratextureid);
    if(cameratextureid==0)
    {
        std::cerr << "Failed to create a texture" << std::endl;
        return false;
    }
    glBindTexture(GL_TEXTURE_2D,cameratextureid);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexImage2D (GL_TEXTURE_2D,0,GL_RGBA,videoframe.cols,videoframe.rows,0,GL_RGBA,GL_UNSIGNED_BYTE,0);
    glBindTexture(GL_TEXTURE_2D,0);


    return true;
}

bool init(void)
{

    glClearColor(0.8f,0.8f,0.8f,1.0f);

    glDisable(GL_LIGHTING);
    glDisable(GL_LIGHT0);

    glDisable(GL_DEPTH_TEST);

    glEnable(GL_TEXTURE_2D);

    glDisable(GL_CULL_FACE);

    glPolygonMode(GL_FRONT_AND_BACK,GL_FILL);

    glPixelStorei(GL_UNPACK_ALIGNMENT,1);
    glPixelStorei(GL_PACK_ALIGNMENT,1);

    if(glsl::checkAvailability()==false)
    {
        std::cerr << "shaders not available" << std::endl;
        return EXIT_FAILURE;
    }

    GLuint shaders[2];
    shaders[0] = glsl::loadVertexShader("../data/shaders/cameraglsl.vert");
    shaders[1] = glsl::loadFragmentShader("../data/shaders/cameraglsl.frag");

    glsl::compile(shaders,2);
    cameraso = glsl::link(shaders,2);
    if(cameraso==0)
        return false;

    createFullScreenQuad(windowsize.width,windowsize.height,cameraQuad);
    glBindVertexArray(cameraQuad.vertexarrayid);
    glBindBuffer(GL_ARRAY_BUFFER, cameraQuad.verticesid);
    cameraQuad.verticesattribid = glGetAttribLocation(cameraso, "vertexPositionIn");
    if(cameraQuad.verticesattribid == -1)
        return false;
    glVertexAttribPointer(cameraQuad.verticesattribid , 3, GL_FLOAT, GL_FALSE, 0, 0);
    glBindBuffer(GL_ARRAY_BUFFER, cameraQuad.texturesid );
    cameraQuad.texturesattribid = glGetAttribLocation(cameraso, "vertexTexCoordIn");
    if(cameraQuad.texturesattribid == -1)
        return false;
    glVertexAttribPointer(cameraQuad.texturesattribid , 2, GL_FLOAT, GL_FALSE, 0, 0);
    glBindVertexArray(0);

    return true;
}


void reshape(int width, int height)
{
    projection = glm::ortho (0.0f, float(windowsize.width), 0.0f, float(windowsize.height), -0.5f, 0.5f);
    mvp = projection;
}


int main(int argc, char **argv)
{

    // Try to open the camera
    cap.open(0);
    if(cap.isOpened()==false)
    {
        hvrl::Log::add().error("main","cannot open the streaming device");
        return EXIT_FAILURE;
    }
    cap.set(CV_CAP_PROP_FRAME_WIDTH,windowsize.width);
    cap.set(CV_CAP_PROP_FRAME_HEIGHT,windowsize.height);
    cap >> videoframe;

    // Init OpenGL
    glutInit(&argc, argv);
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE);
    glutInitWindowPosition(0, 0);
    glutInitWindowSize(windowsize.width, windowsize.height);
    int windowid = -1;
    if((windowid=glutCreateWindow("Undistort with bilinear interpolation")) == 0)
    {
        return EXIT_FAILURE;
    }
    glutShowWindow();
    glutPositionWindow(50, 50);

    glewInit();
    if(!init())
    {
        return EXIT_FAILURE;
    }

    /*
        CUDA
    */
    cuda::checkStatus(cudaGLSetGLDevice(0));
    if(!initCuda())
    {
        return EXIT_FAILURE;
    }

    /*
        Launch the process
    */
    glutReshapeFunc(reshape);
    glutDisplayFunc(display);
    glutKeyboardFunc(keyboard);

    while(isrunning)
    {
        idle();
        glutMainLoopEvent();

    }

    /*
        Cleaning
    */
    glDeleteProgram(cameraso);
    deleteVAO(cameraQuad);
    glDeleteBuffers(1,&pbo);
    glDeleteTextures(1,&cameratextureid);
    cuda::destroyRadialDistortion();
    glutDestroyWindow(windowid);
    cudaGraphicsUnregisterResource(cuda_pbo_resource);

    return 0;
}




bool createFullScreenQuad(const unsigned int& width, const unsigned int& height, VertexArrayObject& vao)
{
    glGenVertexArrays(1, &vao.vertexarrayid);
    if (vao.vertexarrayid == 0)
    {
        hvrl::Log::add().error("createFullScreenQuad","Impossible to generate the vertex array");
        return false;
    }

    vao.verticesattribid = vao.normalsattribid = vao.texturesattribid = -1;

    GLfloat coordinates[12];
    GLfloat textures[8];
    GLuint indices[6];

    textures[0] = textures[1] = textures[3] = textures[6] = 0.0f;
    textures[4] = textures[5] = textures[2] = textures[7] = 1.0f;

    coordinates[0] = coordinates[7] = coordinates[9] = coordinates[10] = 0.0f;
    coordinates[2] = coordinates[5] = coordinates[8] = coordinates[11] = 0.0f;
    coordinates[3] = coordinates[6] = static_cast<float>(width);
    coordinates[1] = coordinates[4] = static_cast<float>(height);

    indices[0] = 0;
    indices[1] = 3;
    indices[2] = 1;
    indices[3] = 1;
    indices[4] = 3;
    indices[5] = 2;

    vao.sizeindices = 6;

    vao.normalsid = 0;
    glGenBuffers(1, &vao.verticesid);
    if (vao.verticesid == 0)
    {
        hvrl::Log::add().error("createFullScreenQuad","Impossible to generate the verticesid");
        return false;
    }
    glGenBuffers(1, &vao.texturesid);
    if (vao.texturesid == 0)
    {
        hvrl::Log::add().error("createFullScreenQuad","Impossible to generate the texturesid");
        return false;
    }
    glGenBuffers(1, &vao.indicesid);
    if (vao.indicesid == 0)
    {
        hvrl::Log::add().error("createFullScreenQuad","Impossible to generate the indicesid");
        return false;
    }

    glBindVertexArray(vao.vertexarrayid);

    glBindBuffer(GL_ARRAY_BUFFER, vao.verticesid);
    glBufferData(GL_ARRAY_BUFFER, 12*sizeof(GLfloat), coordinates, GL_STATIC_DRAW);

    glBindBuffer(GL_ARRAY_BUFFER, vao.texturesid);
    glBufferData(GL_ARRAY_BUFFER, 8*sizeof(GLfloat), textures, GL_STATIC_DRAW);

    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, vao.indicesid);
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, 6*sizeof(GLuint), indices, GL_STATIC_DRAW);

    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, 0);

    return true;

}


void drawVAO(const VertexArrayObject& vao)
{
    if (vao.vertexarrayid != 0)
    {

        glBindVertexArray(vao.vertexarrayid);
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, vao.indicesid);
        if (vao.verticesattribid>=0)
        {
            glEnableVertexAttribArray(vao.verticesattribid);
        }
        if (vao.normalsattribid>=0)
            glEnableVertexAttribArray(vao.normalsattribid);
        if (vao.texturesattribid>=0)
        {
            glEnableVertexAttribArray(vao.texturesattribid);
        }
        glDrawElements(GL_TRIANGLES, vao.sizeindices, GL_UNSIGNED_INT, 0);
        if (vao.verticesattribid>=0)
            glDisableVertexAttribArray(vao.verticesattribid);
        if (vao.normalsattribid>=0)
            glDisableVertexAttribArray(vao.normalsattribid);
        if (vao.texturesattribid>=0)
            glDisableVertexAttribArray(vao.texturesattribid);
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER,0 );
        glBindVertexArray(0);
    }
}

void deleteVAO(VertexArrayObject& vao)
{
    glDeleteVertexArrays(1, &vao.vertexarrayid);
    glDeleteBuffers(1, &vao.verticesid);
    glDeleteBuffers(1, &vao.texturesid);
    glDeleteBuffers(1, &vao.indicesid);
}










